/*
@file: AESMat.cpp
@author: ZZH
@time: 2022-10-20 14:43:38
@info: AES矩阵类, 用于实现数据和密钥
*/

#include "Mat.h"

namespace AES
{
    void Mat_t::setData(const uint8_t* pData, uint8_t len)
    {
        if (len > 16)
            len = 16;

        uint8_t fillLen = 16 - len;

        if (nullptr == pData || 0 == len)
            return;

        memcpy(this->mat, pData, len);

        if (fillLen > 0)
            memset(reinterpret_cast<uint8_t*>(this->mat) + len, fillLen, fillLen);//PKCS7填充规则
    }

    void Mat_t::getData(uint8_t* output, uint8_t len) const
    {
        if (len > 16)
            len = 16;

        if (nullptr == output || 0 == len)
            return;

        memcpy(output, this->mat, len);
    }

    void Mat_t::transpose(void)
    {
        for (int i = 0;i < 3;i++)
        {
            for (int j = i + 1;j < 4;j++)
            {
                uint8_t tmp = this->mat[i][j];
                this->mat[i][j] = this->mat[j][i];
                this->mat[j][i] = tmp;
            }
        }
    }

    void Mat_t::shiftRows(void)
    {
        auto pMat = reinterpret_cast<uint32_t*>(this->mat);

        pMat[1] = ((pMat[1] >> 8) & 0x00FFFFFF) | ((pMat[1] << 24) & 0xFF000000);
        pMat[2] = ((pMat[2] >> 16) & 0x0000FFFF) | ((pMat[2] << 16) & 0xFFFF0000);
        pMat[3] = ((pMat[3] >> 24) & 0x000000FF) | ((pMat[3] << 8) & 0xFFFFFF00);
    }

    void Mat_t::ishiftRows(void)
    {
        auto pMat = reinterpret_cast<uint32_t*>(this->mat);

        pMat[1] = ((pMat[1] << 8) & 0xFFFFFF00) | ((pMat[1] >> 24) & 0x000000FF);
        pMat[2] = ((pMat[2] << 16) & 0xFFFF0000) | ((pMat[2] >> 16) & 0x0000FFFF);
        pMat[3] = ((pMat[3] << 24) & 0xFF000000) | ((pMat[3] >> 8) & 0x00FFFFFF);
    }

    void Mat_t::mix(const uint8_t(*pMixer)[4])
    {
        uint8_t orig[4][4];
        memcpy(orig, this->mat, sizeof(orig));
        for (int i = 0;i < 4;i++)
        {
            for (int j = 0;j < 4;j++)
            {
                this->mat[i][j] =
                    this->mix(pMixer[i][0], orig[0][j]) ^
                    this->mix(pMixer[i][1], orig[1][j]) ^
                    this->mix(pMixer[i][2], orig[2][j]) ^
                    this->mix(pMixer[i][3], orig[3][j]);
            }
        }
    }
} // namespace AES
